package vectorindex_v2

import (
	"context"
	"fmt"
	"time"

	"github.com/cloudwego/eino/schema"
	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
)

// MilvusIndexerConfig 包含 Milvus 索引器的配置
type MilvusIndexerConfig struct {
	// 连接配置
	Address      string
	Username     string
	Password     string
	DatabaseName string

	// 集合配置
	CollectionName string
	VectorDim      int
	MetricType     entity.MetricType // 默认使用 L2 距离

	// 索引配置
	IndexName   string
	IndexType   entity.IndexType
	IndexParams map[string]interface{}

	// 操作配置
	BatchSize int
	Timeout   time.Duration
}

// WithMilvusConfig 创建 Milvus 索引器配置选项
func WithMilvusConfig(config *MilvusIndexerConfig) IndexOption {
	return func(opt *IndexOptions) {
		// 这里简化处理，实际项目中可能需要更复杂的选项传递机制
	}
}

// MilvusIndexer 实现 VectorIndexer 接口
type MilvusIndexer struct {
	client client.Client
	config *MilvusIndexerConfig
}

// NewMilvusIndexer 创建新的 Milvus 索引器
func NewMilvusIndexer(ctx context.Context, config *MilvusIndexerConfig) (VectorIndexer, error) {
	// 验证必需配置
	if config.Address == "" {
		return nil, fmt.Errorf("milvus address is required")
	}
	if config.CollectionName == "" {
		return nil, fmt.Errorf("collection name is required")
	}
	if config.VectorDim <= 0 {
		return nil, fmt.Errorf("vector dimension must be positive")
	}

	// 设置默认值
	if config.MetricType == 0 {
		config.MetricType = entity.L2
	}
	if config.IndexType == "" {
		config.IndexType = entity.IndexHNSW
		config.IndexParams = map[string]interface{}{
			"M":              16,
			"efConstruction": 200,
		}
	}
	if config.BatchSize <= 0 {
		config.BatchSize = 100
	}
	if config.Timeout <= 0 {
		config.Timeout = 30 * time.Second
	}

	// 连接到 Milvus
	client, err := client.NewClient(ctx, client.Config{
		Address:      config.Address,
		Username:     config.Username,
		Password:     config.Password,
		DatabaseName: config.DatabaseName,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to connect to milvus: %w", err)
	}

	// 检查集合是否存在，不存在则创建
	exists, err := client.HasCollection(ctx, config.CollectionName)
	if err != nil {
		return nil, fmt.Errorf("failed to check collection existence: %w", err)
	}

	if !exists {
		// 创建集合
		err = client.CreateCollection(
			ctx,
			config.CollectionName,
			entity.Schema{
				Fields: []entity.Field{
					&entity.Field{
						Name:       "id",
						DataType:   entity.FieldTypeInt64,
						PrimaryKey: true,
						AutoID:     true,
					},
					&entity.Field{
						Name:     "vector",
						DataType: entity.FieldTypeFloatVector,
						TypeParams: map[string]string{
							"dim": fmt.Sprintf("%d", config.VectorDim),
						},
					},
					&entity.Field{
						Name:     "content",
						DataType: entity.FieldTypeVarChar,
						TypeParams: map[string]string{
							"max_length": "65535",
						},
					},
					&entity.Field{
						Name:     "metadata",
						DataType: entity.FieldTypeJSON,
					},
				},
				EnableDynamicField: true,
			},
		)
		if err != nil {
			return nil, fmt.Errorf("failed to create collection: %w", err)
		}

		// 创建索引
		err = client.CreateIndex(
			ctx,
			config.CollectionName,
			"vector",
			entity.IndexType(config.IndexType),
			config.IndexParams,
			false,
		)
		if err != nil {
			return nil, fmt.Errorf("failed to create index: %w", err)
		}
	}

	return &MilvusIndexer{
		client: client,
		config: config,
	}, nil
}

// Store 实现 indexer.Indexer 接口的 Store 方法
func (mi *MilvusIndexer) Store(ctx context.Context, docs []*schema.Document, opts ...IndexOption) ([]string, error) {
	// 解析选项
	commonOpts := GetCommonOptions(nil, opts...)

	// 处理嵌入（如果提供了嵌入器且文档没有向量）
	if commonOpts.Embedding != nil {
		for i := range docs {
			if docs[i].DenseVector() == nil || len(docs[i].DenseVector()) == 0 {
				vectors, err := commonOpts.Embedding.EmbedStrings(ctx, []string{docs[i].Content})
				if err != nil {
					return nil, fmt.Errorf("failed to embed document: %w", err)
				}
				if len(vectors) > 0 {
					docs[i] = docs[i].WithDenseVector(vectors[0])
				}
			}
		}
	}

	// 准备插入数据
	var contents []string
	var vectors []entity.Vector
	var metadatas []string

	for _, doc := range docs {
		// 验证向量维度
		vector := doc.DenseVector()
		if len(vector) != mi.config.VectorDim {
			return nil, fmt.Errorf("vector dimension mismatch: expected %d, got %d", mi.config.VectorDim, len(vector))
		}

		// 转换为 Milvus 向量格式
		floatVector := make([]float32, len(vector))
		for i, v := range vector {
			floatVector[i] = float32(v)
		}

		contents = append(contents, doc.Content)
		vectors = append(vectors, entity.FloatVector(floatVector))

		// 将元数据转换为 JSON 字符串
		metadataJSON, err := doc.MarshalMetaData()
		if err != nil {
			return nil, fmt.Errorf("failed to marshal metadata: %w", err)
		}
		metadatas = append(metadatas, metadataJSON)
	}

	// 批量插入
	resultIDs := make([]string, 0, len(docs))
	for i := 0; i < len(docs); i += mi.config.BatchSize {
		end := i + mi.config.BatchSize
		if end > len(docs) {
			end = len(docs)
		}

		batchContents := contents[i:end]
		batchVectors := vectors[i:end]
		batchMetadatas := metadatas[i:end]

		// 插入数据
		insertResult, err := mi.client.Insert(
			ctx,
			mi.config.CollectionName,
			[]entity.Field{
				entity.NewVarCharField("content", batchContents),
				entity.NewFloatVectorField("vector", mi.config.VectorDim, batchVectors),
				entity.NewJSONField("metadata", batchMetadatas),
			},
		)
		if err != nil {
			return nil, fmt.Errorf("failed to insert data: %w", err)
		}

		// 获取插入的 ID
		for _, id := range insertResult.IDs.Int64IDs {
			resultIDs = append(resultIDs, fmt.Sprintf("%d", id))
		}
	}

	return resultIDs, nil
}

// Close 关闭 Milvus 连接
func (mi *MilvusIndexer) Close() error {
	if mi.client != nil {
		return mi.client.Close()
	}
	return nil
}
